import * as tf from '@tensorflow/tfjs'

function createSimpleModel(index, activation){
  var model = ''
  if(index == 3){
    // 创建3层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[56,56,3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dropout(0.8))
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
    console.log(model);
  }else if(index == 4){
    //创建4层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[56,56,3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dropout(0.8))
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
    console.log(model);
  }else{
    //创建5层全连接网络
    model = tf.sequential();
    model.add(tf.layers.dense({units: 16, inputShape:[56,56,3]}));
    model.add(tf.layers.dense({units: 32, activation: activation}));
    model.add(tf.layers.dense({units: 64, activation: activation}));
    model.add(tf.layers.dropout(0.9))
    model.add(tf.layers.dense({units: 64, activation:  activation}));
    model.add(tf.layers.dropout(0.8))
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
    console.log(model);
  }
  return model;
}

function createComplexModel(index, activation){
  var model = ''
  if(index == 3){
    //创建3层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [56,56,3], kernelSize: 3, strides: 1, filters: 32, activation: activation, padding:'same',kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
  }else if(index == 4){
    //创建4层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [56,56,3], kernelSize: 3, strides: 1, filters: 32, padding:'same', activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 32, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64,activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64,activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
  }else{
    //创建5层卷积神经网络
    model = tf.sequential();
    model.add(tf.layers.conv2d({inputShape: [56,56,3], kernelSize: 3, strides: 1, filters: 32,  padding:'same', activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 32, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 64, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.conv2d({kernelSize: 3, strides: 1, padding:'same', filters: 128, activation: activation,kernelInitializer: 'VarianceScaling'}));
    model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [1, 1]}));
    model.add(tf.layers.flatten());
    model.add(tf.layers.dense({units: 4, activation: 'softmax'}));
  }
  return model;
}

export{
  createSimpleModel,
  createComplexModel
}
